import torch as pt
import numpy as np

# 1.	使用pytorch，完成hello处理（每题10分）
# (1)	数据处理
# ①	将hielo作为字典存储
sentence = 'hihello'
xdict = set(list(sentence))
idx2char = list(xdict)
char2idx = {ch: i for i, ch in enumerate(idx2char)}
N_DICT_LEN = len(idx2char)

# ②	使用hihello获取ihello
x = sentence[:-1]
y = sentence[1:]

# ③	将数据按照要求进行独热处理
x_idx = [char2idx[ch] for ch in x]
y_idx = [char2idx[ch] for ch in y]
x_oh = np.eye(N_DICT_LEN)[x_idx]

# ④	将数据放入tensor
x_idx = pt.tensor([x_idx])
y_idx = pt.tensor([y_idx])
x_oh = pt.tensor([x_oh]).float()

# (2)	模型处理
# ①	设置处理基础参数，类别5输出特征5，隐藏神经5，序列长度6
ALPHA = 0.01
N_CLS = N_DICT_LEN
N_OUTPUT = N_DICT_LEN
N_HIDDEN = 5
N_STEPS = len(x_idx[0])
N_ITERS = 100


# ②	创建模型类
# ③	使用rnn对模型进行处理
# ④	创建正向传播
class MyRnn(pt.nn.Module):

    def __init__(self, input_size, hidden_size, num_layers, n_cls, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.rnn = pt.nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = pt.nn.Linear(hidden_size, n_cls)

    def forward(self, x):
        x, _ = self.rnn(x)
        x = x.reshape(-1, self.hidden_size)
        x = self.fc(x)
        return x


model = MyRnn(N_DICT_LEN, N_HIDDEN, 1, N_DICT_LEN)
criterion = pt.nn.CrossEntropyLoss()
optim = pt.optim.Adam(params=model.parameters(), lr=ALPHA)


def acc(h, y):
    return h.argmax(1).eq(y.long()).double().mean()


# ⑤	模型编译测试，反向传播
# ⑥	打印训练后的预测结果
GROUP = int(np.ceil(N_ITERS / 20))
y_idx = y_idx.reshape(-1)
for i in range(N_ITERS):
    model.train(True)
    optim.zero_grad()
    h_oh = model(x_oh)
    cost = criterion(h_oh, y_idx)
    cost.backward()
    optim.step()
    model.train(False)
    cost = cost.data.numpy()
    accv = acc(h_oh, y_idx).data.numpy()
    h = h_oh.argmax(1).reshape(-1, N_STEPS)
    h = h.data.numpy()
    str = [''.join([idx2char[idx] for idx in hi]) for hi in h]
    if i % GROUP == 0:
        print(f'#{i + 1}, cost = {cost}, acc = {accv}, prediction = {str}')
if i % GROUP != 0:
    print(f'#{i + 1}, cost = {cost}, acc = {accv}, prediction = {str}')
